# -*- coding: UTF-8 -*-
# @IDE     : VScode
# @File   : chat_bot_serve.py
# @Time   : 2024/05/25 15:35:39
# @Author : zhonggc

from dotenv import load_dotenv, find_dotenv
import os
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_cohere import ChatCohere
from langchain_core.output_parsers import StrOutputParser
from fastapi import FastAPI
from langserve import add_routes
import uvicorn
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables import RunnablePassthrough
from langchain_core.messages import HumanMessage
from LLM.zhipu_llm import ZhiPu


app = FastAPI(
    title="langChain chatBot Server",
    version="1.0",
    description="A simple chatBot API server",
)


class ChatBotServe:

    def __init__(self):
        self.setup()
        self.store = {}
        self.messages = []

    def setup(self):
        # 加载环境变量
        _ = load_dotenv(find_dotenv())
        self.COHERE_API_KEY = os.getenv("COHERE_API_KEY")

    def create_prompt_template(self):
        """构建对话模板
        """
        self.prompt = ChatPromptTemplate.from_messages(
            [
                ("system", "你是一个聊天机器人，你的任务是尽可能详细并准确地回答用户的问题。"),
                MessagesPlaceholder(variable_name="messages"),
            ],
        )

    def create_model(self):
        """构建对话模型
        """
        # self.model = ChatCohere(model="command-r")
        ZHIPUAI_API_KEY = os.getenv("ZHIPUAI_API_KEY")
        self.model = ZhiPu(api_key=ZHIPUAI_API_KEY)

    def get_session_history(self, session_id):
        """获取会话历史记录
        """
        if session_id not in self.store:
            self.store[session_id] = ChatMessageHistory()
        return self.store.get(session_id, None)

    def create_chain(self):
        """构建对话链
        """
        def filter_messages(messages, k=10):
            if len(messages) > k:
                self.messages = messages[-k:]
            else:
                self.messages = messages
            return self.messages

        parser = StrOutputParser()
        self.chain = (
            RunnablePassthrough.assign(
                messages=lambda x: filter_messages(x['messages']))
            | self.prompt
            | self.model
            | parser
        )

    def run(self):
        self.create_prompt_template()
        self.create_model()
        self.create_chain()
        add_routes(
            app,
            self.chain,
            path="/chain",
        )
        uvicorn.run(app, host="localhost", port=8000)

chat_bot_serve = ChatBotServe()

@app.post("/get_response/")
def get_response(data: dict):
    """获取回复
    """
    text = data.get("text", "")
    with_message_history = RunnableWithMessageHistory(
        chat_bot_serve.chain,
        chat_bot_serve.get_session_history,
        input_messages_key="messages",
    )
    config = {
        "configurable": {
            "session_id": "123",
        }
    }

    response = with_message_history.invoke(
        {
            "messages": chat_bot_serve.messages + [HumanMessage(content=text)]
        },
        config=config,
    )
    return {"response": response, "text": text}


if __name__ == "__main__":
    chat_bot_serve.run()
